{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo of simple density-based imitation learning baselines\n",
    "\n",
    "This demo shows how to train a `Pendulum` agent (exciting!) with our simple density-based imitation learning baselines. `DensityTrainer` has a few interesting parameters, but the key ones are:\n",
    "\n",
    "1. `density_type`: this governs whether density is measured on $(s,s')$ pairs (`db.STATE_STATE_DENSITY`), $(s,a)$ pairs (`db.STATE_ACTION_DENSITY`), or single states (`db.STATE_DENSITY`).\n",
    "2. `is_stationary`: determines whether a separate density model is used for each time step $t$ (`False`), or the same model is used for transitions at all times (`True`).\n",
    "3. `standardise_inputs`: if `True`, each dimension of the agent state vectors will be normalised to have zero mean and unit variance over the training dataset. This can be useful when not all elements of the demonstration vector are on the same scale, or when some elements have too wide a variation to be captured by the fixed kernel width (1 for Gaussian kernel).\n",
    "4. `kernel`: changes the kernel used for non-parametric density estimation. `gaussian` and `exponential` are the best bets; see the [sklearn docs](https://scikit-learn.org/stable/modules/density.html#kernel-density) for the rest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "#%load_ext autoreload\n",
    "#%autoreload 2\n",
    "\n",
    "import pprint\n",
    "\n",
    "from imitation.algorithms import density_baselines as db\n",
    "from imitation.data import types\n",
    "from imitation.util import util\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Set FAST = False for longer training. Use True for testing and CI.\n",
    "FAST = True\n",
    "\n",
    "if FAST:\n",
    "    N_VEC = 1\n",
    "    N_TRAJECTORIES = 1\n",
    "    N_ITERATIONS = 1\n",
    "    N_RL_TRAIN_STEPS = 100\n",
    "else:\n",
    "    N_VEC = 8\n",
    "    N_TRAJECTORIES = 10\n",
    "    N_ITERATIONS = 100\n",
    "    N_RL_TRAIN_STEPS = int(1e5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n"
     ]
    }
   ],
   "source": [
    "env_name = 'Pendulum-v0'\n",
    "env = util.make_vec_env(env_name, N_VEC)\n",
    "rollouts = types.load(\"../tests/testdata/expert_models/pendulum_0/rollouts/final.pkl\")\n",
    "imitation_trainer = util.init_rl(env, learning_rate=3e-4, n_steps=2048)\n",
    "density_trainer = db.DensityTrainer(env,\n",
    "                                    rollouts=rollouts,\n",
    "                                    imitation_trainer=imitation_trainer,\n",
    "                                    density_type=db.STATE_ACTION_DENSITY,\n",
    "                                    is_stationary=False,\n",
    "                                    kernel='gaussian',\n",
    "                                    kernel_bandwidth=0.2,  # found using divination & some palm reading\n",
    "                                    standardise_inputs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Novice stats (true reward function):\n",
      "{'len_max': 200,\n",
      " 'len_mean': 200.0,\n",
      " 'len_min': 200,\n",
      " 'len_std': 0.0,\n",
      " 'monitor_return_max': -790.08163,\n",
      " 'monitor_return_mean': -1219.5085344000001,\n",
      " 'monitor_return_min': -1640.48422,\n",
      " 'monitor_return_std': 264.06817071118337,\n",
      " 'n_traj': 10,\n",
      " 'return_max': -790.0816296990961,\n",
      " 'return_mean': -1219.5085344826803,\n",
      " 'return_min': -1640.4842166900635,\n",
      " 'return_std': 264.06817043374}\n",
      "Novice stats (imitation reward function):\n",
      "{'len_max': 200,\n",
      " 'len_mean': 200.0,\n",
      " 'len_min': 200,\n",
      " 'len_std': 0.0,\n",
      " 'monitor_return_max': -1356.732713,\n",
      " 'monitor_return_mean': -1356.732713,\n",
      " 'monitor_return_min': -1356.732713,\n",
      " 'monitor_return_std': 0.0,\n",
      " 'n_traj': 1,\n",
      " 'return_max': -3049.44188952446,\n",
      " 'return_mean': -3049.44188952446,\n",
      " 'return_min': -3049.44188952446,\n",
      " 'return_std': 0.0}\n",
      "----------------------------------\n",
      "| rollout/           |           |\n",
      "|    ep_len_mean     | 200       |\n",
      "|    ep_rew_mean     | -1.27e+03 |\n",
      "| time/              |           |\n",
      "|    fps             | 284       |\n",
      "|    iterations      | 1         |\n",
      "|    time_elapsed    | 7         |\n",
      "|    total_timesteps | 2048      |\n",
      "----------------------------------\n",
      "Trained stats (epoch 0):\n",
      "{'len_max': 200,\n",
      " 'len_mean': 200.0,\n",
      " 'len_min': 200,\n",
      " 'len_std': 0.0,\n",
      " 'monitor_return_max': -1704.242057,\n",
      " 'monitor_return_mean': -1704.242057,\n",
      " 'monitor_return_min': -1704.242057,\n",
      " 'monitor_return_std': 0.0,\n",
      " 'n_traj': 1,\n",
      " 'return_max': -1704.2420597076416,\n",
      " 'return_mean': -1704.2420597076416,\n",
      " 'return_min': -1704.2420597076416,\n",
      " 'return_std': 0.0}\n",
      "Trained stats (imitation reward function, epoch 0):\n",
      "{'len_max': 200,\n",
      " 'len_mean': 200.0,\n",
      " 'len_min': 200,\n",
      " 'len_std': 0.0,\n",
      " 'monitor_return_max': -890.187322,\n",
      " 'monitor_return_mean': -1150.8099041,\n",
      " 'monitor_return_min': -1386.683418,\n",
      " 'monitor_return_std': 146.3703252135178,\n",
      " 'n_traj': 10,\n",
      " 'return_max': -2667.4581129550934,\n",
      " 'return_mean': -2950.378186058998,\n",
      " 'return_min': -3491.6277445554733,\n",
      " 'return_std': 228.27342110482073}\n"
     ]
    }
   ],
   "source": [
    "novice_stats = density_trainer.test_policy()\n",
    "print('Novice stats (true reward function):')\n",
    "pprint.pprint(novice_stats)\n",
    "novice_stats_im = density_trainer.test_policy(true_reward=False, n_trajectories=N_TRAJECTORIES)\n",
    "print('Novice stats (imitation reward function):')\n",
    "pprint.pprint(novice_stats_im)\n",
    "\n",
    "for i in range(N_ITERATIONS):\n",
    "    density_trainer.train_policy(N_RL_TRAIN_STEPS)\n",
    "\n",
    "    good_stats = density_trainer.test_policy(n_trajectories=N_TRAJECTORIES)\n",
    "    print(f'Trained stats (epoch {i}):')\n",
    "    pprint.pprint(good_stats)\n",
    "    novice_stats_im = density_trainer.test_policy(true_reward=False)\n",
    "    print(f'Trained stats (imitation reward function, epoch {i}):')\n",
    "    pprint.pprint(novice_stats_im)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
